﻿//-------------------------------------
// WBFSSync - WBFSSync.exe
//
// Copyright 2009 Caian (ÔmΣga Frøst) <frost.omega@hotmail.com>
//
// WBFSSync is Licensed under the terms of the Microsoft Reciprocal License (Ms-RL)
//
// DeviceManager.cs:
//
// Implementa gerenciador de drives
//
//-------------------------------------

using System;
using System.Collections.Generic;
using System.Linq;
using System.Text;
using System.Runtime.InteropServices;
using System.Collections;
using System.IO;
using System.Threading;
using System.Windows.Forms;
using System.Management;

namespace WBFSSync
{
    public delegate void DeviceListChangedDelegate(bool devicesInserted, bool devicesRemoved, bool devicesUpdated);

    public static class DeviceManager
    {
        //---------------------------- Estruturas internas úteis

        struct DEV_BROADCAST_HDR //Estrutura que carrega informação do tipo de dispositivo inserido/removido
        {
            public uint dbch_size;
            public uint dbch_devicetype;
            public uint dbch_reserved;

            public DEV_BROADCAST_HDR(uint s, uint d) // Só pro compilador parar de dar Warning
            {
                dbch_size = s;
                dbch_devicetype = d;
                dbch_reserved = 0;
            }
        }

        struct DEV_BROADCAST_VOLUME //Estrutura que carrega informação sobre os volume(s) inserido(s)/removido(s)
        {
            public uint dbcv_size;
            public uint dbcv_devicetype;
            public uint dbcv_reserved;
            public uint dbcv_unitmask;
            public ushort dbcv_flags;

            public DEV_BROADCAST_VOLUME(uint s, uint d, uint m, ushort f) // Só pro compilador parar de dar Warning
            {
                dbcv_size = s;
                dbcv_devicetype = d;
                dbcv_reserved = 0;
                dbcv_unitmask = m;
                dbcv_flags = f;
            }
        }

        //---------------------------- Constantes

        const int WM_DEVICECHANGE = 537; //Mensagem de inserção/remoção de dispositivo USB

        const int DBT_DEVICEARRIVAL = 32768; //WParam para dispositivo inserido
        const int DBT_DEVICEREMOVECOMPLETE = 32772; //WParam para dispositivo removido

        const uint DBT_DEVTYP_VOLUME = 0x00000002; //Dispositivo do tipo 'Volume'

        const uint WM_GETTEXT = 13;
        const uint GW_HWNDNEXT = 2;
        const uint GW_CHILD = 5;
        const uint BM_CLICK = 245;

        //---------------------------- Variáveis

        static List<IDrive> WindowsDrives = new List<IDrive>();
        static List<IDrive> WBFSDrives = new List<IDrive>();

        static Queue<String> PendingInsertions = new Queue<String>();
        static Queue<String> PendingRemotions = new Queue<String>();

        public static Int32 DriveCount { get { return WindowsDrives.Count + WBFSDrives.Count; } }
        public static Int32 WindowsDriveCount { get { return WindowsDrives.Count; } }
        public static Int32 WBFSDriveCount { get { return WBFSDrives.Count; } }

        public static String FormatLockedDrives = ""; //Drives que não podem ser formatados a pedido do usuário
        public static String IgnoredDrives = "AB"; //Drives serão ignorados pelo gerenciador

        public static bool SupressFormatDriveWindow = true;
        
        public static string FormatDriveWindowTitle = "Microsoft Windows";
        public static string FormatDriveWindowClass = "#32770";

        public const int IDCANCEL = 2;
        public static string FormatDriveButtonClass = "Button";
        public static string FormatDriveCancelMessage = "Cancelar";
        //public static string FormatDriveFormatMessage = "Formatar disco";

        public static event DeviceListChangedDelegate DeviceListChanged;

        //---------------------------- Rotinas

        //--------------- Mapeia todas as unidades
        public static void MapAllDevices()
        {
            DriveInfo[] drives = DriveInfo.GetDrives();
            foreach (DriveInfo d in drives) PendingInsertions.Enqueue(d.Name);

            ProcessPendingDrives();
        }

        //--------------- Fecha todas as unidades
        public static void CloseAllDevices()
        {
            foreach (IDrive d in WindowsDrives) PendingRemotions.Enqueue(d.Name);
            foreach (IDrive d in WBFSDrives) PendingRemotions.Enqueue(d.Name);

            ProcessPendingDrives(false);
        }

        //--------------- Mapeia todas as unidades
        public static void AddVirtualDevice(String path)
        {
            path = Path.GetFullPath(path);
            if (!File.Exists(path)) return;
            PendingInsertions.Enqueue(path);
            ProcessPendingDrives();
        }

        //--------------- Adiciona um drive para ser removido na fila
        public static void EnqueueDriveRemotion(string name)
        {
            PendingRemotions.Enqueue(name);
        }

        //--------------- Adiciona um drive para ser inserido na fila
        public static void EnqueueDriveInsertions(string name)
        {
            PendingInsertions.Enqueue(name);
        }

        //--------------- Processa o fechamento e a abertura dos drives pendentes
        public static void ProcessPendingDrives()
        {
            ProcessPendingDrives(true);
        }

        public static void ProcessPendingDrives(bool checkBeforeRemove)
        {
            bool devicesremoved = false;
            bool devicesinserted = false;
            bool devicesupdated = false;

            while (PendingRemotions.Count > 0)
            {
                String d = PendingRemotions.Dequeue();
                IDrive drive = GetDrive(d);

                if (drive != null) // drive == null Não deveria acontecer em condições normais
                {
                    //Verifica se o drive foi desconectado ou, como no caso de leitores de CD, apenas a
                    //mídia foi removida
                    if (!PhysicalDriveExists(drive.Name) || !checkBeforeRemove)
                    {
                        if (drive.Is(DriveFlags.WBFS))
                            WBFSDrives.Remove(drive);
                        else
                            WindowsDrives.Remove(drive);
                    }

                    //Fecha o drive
                    drive.Close();

                    devicesremoved = true;
                }
            }

            while (PendingInsertions.Count > 0)
            {
                String d = PendingInsertions.Dequeue();
                if (IgnoredDrives.Contains(Char.ToUpper(d[0]))) continue;

                IDrive drive = GetDrive(d);
                if (drive == null) // drive != null só deveria acontecer com drives de CD/DVD/BD/Disquete
                {
                    DriveFlags flags = DriveFlags.Unknown;
                    DriveInfo info = new DriveInfo(d.ToString());

                    switch (info.DriveType)
                    {
                        case DriveType.CDRom:
                            flags |= DriveFlags.CDRom;
                            break;
                        case DriveType.Fixed:
                            flags |= DriveFlags.Fixed;
                            break;
                        case DriveType.Network:
                            flags |= DriveFlags.Network;
                            break;
                        case DriveType.Removable:
                            flags |= DriveFlags.Removable;
                            break;

                        default:
                            continue; //Não é um drive válido
                    }

                    Boolean isWBFS = false;
                    if ((info.DriveType == DriveType.Fixed) || (info.DriveType == DriveType.Removable)) 
                        isWBFS = IsWBFSDrive(d);

                    if (isWBFS)
                    {
                        drive = new WBFSDrive();
                        (drive as WBFSDrive)._Device.Open(d, false);
                        (drive as WBFSDrive)._Flags = flags | DriveFlags.WBFS;

                        if ((drive as WBFSDrive)._Device.Virtual)
                            (drive as WBFSDrive)._Flags |= DriveFlags.Virtual;

                        WBFSDrives.Add(drive);
                    }
                    else
                    {
                        if ((flags & DriveFlags.CDRom) != 0)
                        {
                            drive = new CDDrive();

                            (drive as CDDrive)._Drive = info;
                            (drive as CDDrive)._Flags = flags;

                            drive.Refresh();
                        }
                        else
                        {
                            drive = new WinDrive();

                            if (Char.ToLower(Environment.SystemDirectory[0]) == Char.ToLower(d[0]))
                                flags |= DriveFlags.WinDrive;

                            (drive as WinDrive)._Drive = info;
                            (drive as WinDrive)._Flags = flags;
                        }

                        WindowsDrives.Add(drive);
                    }

                    devicesinserted = true;
                }
                else
                {
                    drive.Refresh();
                    devicesupdated = true;
                }
            }

            OnDeviceListChanged(devicesinserted, devicesremoved, devicesupdated);
        }

        //--------------- Atualiza um drive e ativa o evento OnDeviceListChanged(), apenas chamar IDrive.Refresh()
        //não garante que o programa principal atualize a sua exibição, se necessário.
        public static void UpdateDrive(IDrive drive)
        {
            drive.Refresh();
            OnDeviceListChanged(false, false, true);
        }

        //--------------- Verifica se o drive físico ainda existe, não se aplica para dispositivos virtuais
        public static Boolean PhysicalDriveExists(String name)
        {
            String[] drives = Environment.GetLogicalDrives();
            return drives.Contains(name);
        }

        //--------------- Retorna o caminho de um drive de rede
        public static string GetDrivePath(String drive)
        {
            if (drive.EndsWith("\\")) 
                drive = drive.Remove(drive.Length - 1);

            ObjectQuery oq = new ObjectQuery("select RemotePath from Win32_NetworkConnection where LocalName = \"" + drive + '"');
            ManagementScope scope = new ManagementScope("\\root\\cimv2");
            scope.Connect();

            ManagementObjectSearcher sea = new ManagementObjectSearcher(scope, oq);
            ManagementObjectCollection moc = sea.Get();

            foreach (ManagementObject proc in moc) 
                return proc["RemotePath"].ToString();

            return String.Empty;
        }

        //--------------- Verifica se o drive é WBFS
        public static Boolean IsWBFSDrive(String c)
        {
            foreach (IDrive d in WBFSDrives) if (String.Compare(c, d.Name, true) == 0) return true;

            //O Drive ainda não foi aberto, faz a verificação pelo cabeçalho do disco
            int r = WBFSDevice.IsWBFSDrive(c, false);

            if (r == (int)WBFSRet.RET_WBFS_OK) // É um drive WBFS
            {
                return true;
            }
            else if ((r == (int)WBFSRet.RET_WBFS_BADFSMAGIC) || //Não é um drive WBFS
                     (r == (int)WBFSRet.RET_WBFS_BADFSSECSZ) ||
                     (r == (int)WBFSRet.RET_WBFS_BADFSTSECS))
            {
                return false;
            }
            else //r < 0 Erro durante a abertura do disco, verificar por acesso negado
            {    //Acesso negado não deveria ocorrer uma vez que no Vista/Seven ele força rodar como Administrador
                return false;
            }
        }

        //--------------- Verifica se o drive é WBFS
        public static Boolean IsWBFSDrive(int i)
        {
            if ((i < WindowsDrives.Count) || (i >= WindowsDrives.Count + WBFSDrives.Count)) 
                return false;
            else 
                return true;
        }

        //--------------- Verifica se o drive é virtual, não é para nenhum WinDrive conter DriveFlags.Virtual
        public static Boolean IsVirtualDrive(int i)
        {
            if (i < 0) return false;
            else if (i < WindowsDrives.Count) return WindowsDrives[i].Is(DriveFlags.Virtual);
            else if (i < WindowsDrives.Count + WBFSDrives.Count) return WBFSDrives[i - WindowsDrives.Count].Is(DriveFlags.Virtual);
            else return false;
        }

        //--------------- Retorna o drive correspondente à letra selecionada
        public static IDrive GetDrive(String c)
        {
            foreach (IDrive d in WindowsDrives) if (String.Compare(c, d.Name, true) == 0) return d;
            foreach (IDrive d in WBFSDrives) if (String.Compare(c, d.Name, true) == 0) return d;
            return null;
        }

        //--------------- Retorna o drive correspondente ao indice selecionado, drives WBFS por ultimo
        public static IDrive GetDrive(int i)
        {
            if (i < 0) return null;
            else if (i < WindowsDrives.Count) return WindowsDrives[i];
            else if (i < WindowsDrives.Count + WBFSDrives.Count) return WBFSDrives[i - WindowsDrives.Count];
            else return null;
        }

        //--------------- Retorna o drive de indice relativo a uma lista de IDrives que possuem um certa flag
        public static IDrive GetDrive(int i, DriveFlags flags)
        {
            int j = 0;
            for(int k = 0; k < DriveCount; i++)
            {
                IDrive d = GetDrive(i);
                if (d == null) continue;

                if (d.Is(flags))
                {
                    if (i == j) return d;
                    else j++;
                }
            }

            return null;
        }

        //--------------- Retorna o drive não wbfs correspondente ao indice selecionado
        public static IDrive GetWindowsDrive(int i)
        {
            if (i < 0) return null;
            else if (i < WindowsDrives.Count) return WindowsDrives[i];
            else return null;
        }

        //--------------- Retorna o drive wbfs correspondente ao indice selecionado
        public static IDrive GetWBFSDrive(int i)
        {
            if (i < 0) return null;
            else if (i < WBFSDrives.Count) return WBFSDrives[i];
            else return null;
        }

        //--------------- Processa mensagens de adição / remoção de dispositivos
        public static void ProcessSystemMessage(int msg, long wparam, IntPtr lparam)
        {
            //É uma inserção/remoção de um dispositivo?
            if ((msg == WM_DEVICECHANGE) && ((wparam == DBT_DEVICEARRIVAL) || (wparam == DBT_DEVICEREMOVECOMPLETE)))
            {
                Boolean IsArrival = (wparam == DBT_DEVICEARRIVAL);

                //Adquire a estrutura que representa o dispositivo
                DEV_BROADCAST_HDR device = new DEV_BROADCAST_HDR();
                device = (DEV_BROADCAST_HDR)Marshal.PtrToStructure(lparam, typeof(DEV_BROADCAST_HDR));

                //É um volume?
                if (device.dbch_devicetype == DBT_DEVTYP_VOLUME)
                {
                    DEV_BROADCAST_VOLUME volume = new DEV_BROADCAST_VOLUME();
                    volume = (DEV_BROADCAST_VOLUME)Marshal.PtrToStructure(lparam, typeof(DEV_BROADCAST_VOLUME));

                    Queue<String> list = (IsArrival ? PendingInsertions : PendingRemotions);

                    //Adquire lista dos drives modificados
                    for (int i = 0; i < 32; i++)
                    {
                        if ((volume.dbcv_unitmask & (1 << i)) != 0) list.Enqueue(((Char)('A' + i)).ToString() + ":\\");
                    }

                    list = null;

                    Thread.Sleep(500);
                    if (IsArrival && SupressFormatDriveWindow) BeginSupressFormatDrivePopUp();
                    ProcessPendingDrives();
                }
            }
        }

        //--------------- Acha qual texto o windows está usando para o botão 'Cancelar'
        public static string GetCancelString()
        {
            const int maxloop = 100;
            int maxiter = 1000;
            int window = 0;
            int i = 0;

            Thread dialogthread = ThreadManager.LaunchThread("DevMan.ShwTstDlg", new ThreadStart(ShowTestDialog));

            for (i = 0; i < maxloop; i++)
            {
                if ((window = FindWindow("#32770", "WBFSSYNCCT")) != 0) break;
                Thread.Sleep(10);
            }

            if (i == maxloop)
            {
                dialogthread.Abort();
                return String.Empty;
            }

            String canceltext = RecursiveGetIdCancel(window, ref maxiter);
            if (canceltext == String.Empty)
            {
                dialogthread.Abort();
                return String.Empty;
            }

            return canceltext;
        }

        //--------------- Thread da rotina acima
        static void ShowTestDialog()
        {
            MessageBox.Show("CANCEL BUTTON TEST", "WBFSSYNCCT", MessageBoxButtons.OKCancel);
        }

        //--------------- Recursiva para achar o IDCANCEL
        static string RecursiveGetIdCancel(int hWnd, ref int counter)
        {
            if (--counter < 0) //Previne Possíveis Loops gigantes
            {
                return String.Empty;
            }

            StringBuilder classname = new StringBuilder(256);
            RealGetWindowClass(hWnd, classname, 256);

            if (classname.ToString() == FormatDriveButtonClass)
            {
                int id = GetDlgCtrlID(hWnd);
                if (id == IDCANCEL)
                {
                    StringBuilder text = new StringBuilder(256);
                    SendMessage(hWnd, WM_GETTEXT, 256, text);
                    SendMessage(hWnd, BM_CLICK, 0, 0); //Fecha o dialogo
                    return text.ToString();
                }
            }

            String s;
            for (int childWnd = GetWindow(hWnd, GW_CHILD); childWnd != 0; childWnd = GetWindow(childWnd, GW_HWNDNEXT))
            {
                if (counter <= 0) return String.Empty;
                if ((s = RecursiveGetIdCancel(childWnd, ref counter)) != String.Empty) return s;
            }

            return String.Empty;
        }

        //--------------- Cancela a janela de "Formatar Disco"
        public static void BeginSupressFormatDrivePopUp()
        {
            ThreadManager.LaunchThread("DevMan.SprssFrmt", new ThreadStart(SupressFormatDrivePopUp));
        }

        //--------------- Cancela a janela de "Formatar Disco"
        public static void SupressFormatDrivePopUp()
        {
            int hWnd;

            StringBuilder classname = new StringBuilder(256);
            StringBuilder text = new StringBuilder(256);

            //Procura pela janela
            long before = DateTime.Now.Ticks;
            while ((hWnd = FindWindow(FormatDriveWindowClass, FormatDriveWindowTitle)) == 0)
            {
                if (TimeSpan.FromTicks(DateTime.Now.Ticks - before).TotalSeconds > 5)
                    return;

                Thread.Sleep(1);
            }

            Thread.Sleep(50);

            int maxiter = 100; //Limita o numero de interações dentro dos loops
            RecursiveCancelFormat(hWnd, ref maxiter);
        }

        public static bool RecursiveCancelFormat(int hWnd, ref int counter)
        {
            if (--counter < 0) //Previne Possíveis Loops gigantes
            {
                return false;
            }

            StringBuilder classname = new StringBuilder(256);
            RealGetWindowClass(hWnd, classname, 256);

            if (classname.ToString() == FormatDriveButtonClass)
            {
                StringBuilder text = new StringBuilder(256);
                SendMessage(hWnd, WM_GETTEXT, 256, text);

                if (text.ToString() == FormatDriveCancelMessage)
                {
                    SendMessage(hWnd, BM_CLICK, 0, 0);
                    return true;
                }
            }

            for (int childWnd = GetWindow(hWnd, GW_CHILD); childWnd != 0; childWnd = GetWindow(childWnd, GW_HWNDNEXT))
            {
                if (counter <= 0) return false;
                if (RecursiveCancelFormat(childWnd, ref counter)) return true;
            }

            return false;
        }

        //--------------- Dispara o evento de mudança de drives
        public static void OnDeviceListChanged(bool devicesInserted, bool devicesRemoved, bool devicesUpdated)
        {
            if (DeviceListChanged != null) DeviceListChanged(devicesInserted, devicesRemoved, devicesUpdated);
        }

        //--------------- Rotinas nativas
        [DllImport("user32.dll")]
        private static extern int SendMessage(int hWnd, uint message, int size, StringBuilder text);
        [DllImport("user32.dll")]
        private static extern int SendMessage(int hWnd, uint message, uint wparam, long lparam);
        [DllImport("user32.dll")]
        private static extern int GetWindow(int hWnd, uint uCmd);
        [DllImport("user32.dll")]
        private static extern int RealGetWindowClass(int hWnd, StringBuilder title, int size);
        [DllImport("user32.dll")]
        private static extern int FindWindow(string lpClassName, string lpWindowName);
        [DllImport("user32.dll")]
        static extern int GetDlgCtrlID(int hwndCtl);
    }
}
